{
 "metadata": {
  "signature": "sha256:30840350b63c24edb7cccc0975926b54a69bfbab248af75223666466c52cf8af"
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "The FortranFile class\n",
      "=====================\n",
      "\n",
      "> **NOTE**: you may want to use [scipy.io.FortranFile](http://docs.scipy.org/doc/scipy/reference/generated/scipy.io.FortranFile.html) instead.\n",
      "\n",
      "This subclass of file is designed to simplify reading of Fortran\n",
      "unformatted binary files which are typically saved in a record-based\n",
      "format."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# Copyright 2008, 2009 Neil Martinsen-Burrell\n",
      "#\n",
      "# Permission is hereby granted, free of charge, to any person obtaining a copy\n",
      "# of this software and associated documentation files (the \"Software\"), to deal\n",
      "# in the Software without restriction, including without limitation the rights\n",
      "# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n",
      "# copies of the Software, and to permit persons to whom the Software is\n",
      "# furnished to do so, subject to the following conditions:\n",
      "\n",
      "# The above copyright notice and this permission notice shall be included in\n",
      "# all copies or substantial portions of the Software.\n",
      "\n",
      "# THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n",
      "# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n",
      "# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n",
      "# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n",
      "# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n",
      "# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN\n",
      "# THE SOFTWARE.\n",
      "\n",
      "\"\"\"Defines a file-derived class to read/write Fortran unformatted files.\n",
      "\n",
      "The assumption is that a Fortran unformatted file is being written by\n",
      "the Fortran runtime as a sequence of records.  Each record consists of\n",
      "an integer (of the default size [usually 32 or 64 bits]) giving the\n",
      "length of the following data in bytes, then the data itself, then the\n",
      "same integer as before.\n",
      "\n",
      "Examples\n",
      "--------\n",
      "\n",
      "To use the default endian and size settings, one can just do::\n",
      "    >>> f = FortranFile('filename')\n",
      "    >>> x = f.readReals()\n",
      "\n",
      "One can read arrays with varying precisions::\n",
      "    >>> f = FortranFile('filename')\n",
      "    >>> x = f.readInts('h')\n",
      "    >>> y = f.readInts('q')\n",
      "    >>> z = f.readReals('f')\n",
      "Where the format codes are those used by Python's struct module.\n",
      "\n",
      "One can change the default endian-ness and header precision::\n",
      "    >>> f = FortranFile('filename', endian='>', header_prec='l')\n",
      "for a file with little-endian data whose record headers are long\n",
      "integers.\n",
      "\"\"\"\n",
      "\n",
      "__docformat__ = \"restructuredtext en\"\n",
      "\n",
      "import struct\n",
      "import numpy\n",
      "\n",
      "class FortranFile(file):\n",
      "\n",
      "    \"\"\"File with methods for dealing with fortran unformatted data files\"\"\"\n",
      "\n",
      "    def _get_header_length(self):\n",
      "        return struct.calcsize(self._header_prec)\n",
      "    _header_length = property(fget=_get_header_length)\n",
      "\n",
      "    def _set_endian(self,c):\n",
      "        \"\"\"Set endian to big (c='>') or little (c='<') or native (c='@')\n",
      "\n",
      "        :Parameters:\n",
      "          `c` : string\n",
      "            The endian-ness to use when reading from this file.\n",
      "        \"\"\"\n",
      "        if c in '<>@=':\n",
      "            self._endian = c\n",
      "        else:\n",
      "            raise ValueError('Cannot set endian-ness')\n",
      "    def _get_endian(self):\n",
      "        return self._endian\n",
      "    ENDIAN = property(fset=_set_endian,\n",
      "                      fget=_get_endian,\n",
      "                      doc=\"Possible endian values are '<', '>', '@', '='\"\n",
      "                     )\n",
      "\n",
      "    def _set_header_prec(self, prec):\n",
      "        if prec in 'hilq':\n",
      "            self._header_prec = prec\n",
      "        else:\n",
      "            raise ValueError('Cannot set header precision')\n",
      "    def _get_header_prec(self):\n",
      "        return self._header_prec\n",
      "    HEADER_PREC = property(fset=_set_header_prec,\n",
      "                           fget=_get_header_prec,\n",
      "                           doc=\"Possible header precisions are 'h', 'i', 'l', 'q'\"\n",
      "                          )\n",
      "\n",
      "    def __init__(self, fname, endian='@', header_prec='i', *args, **kwargs):\n",
      "        \"\"\"Open a Fortran unformatted file for writing.\n",
      "        \n",
      "        Parameters\n",
      "        ----------\n",
      "        endian : character, optional\n",
      "            Specify the endian-ness of the file.  Possible values are\n",
      "            '>', '<', '@' and '='.  See the documentation of Python's\n",
      "            struct module for their meanings.  The deafult is '>' (native\n",
      "            byte order)\n",
      "        header_prec : character, optional\n",
      "            Specify the precision used for the record headers.  Possible\n",
      "            values are 'h', 'i', 'l' and 'q' with their meanings from\n",
      "            Python's struct module.  The default is 'i' (the system's\n",
      "            default integer).\n",
      "\n",
      "        \"\"\"\n",
      "        file.__init__(self, fname, *args, **kwargs)\n",
      "        self.ENDIAN = endian\n",
      "        self.HEADER_PREC = header_prec\n",
      "\n",
      "    def _read_exactly(self, num_bytes):\n",
      "        \"\"\"Read in exactly num_bytes, raising an error if it can't be done.\"\"\"\n",
      "        data = ''\n",
      "        while True:\n",
      "            l = len(data)\n",
      "            if l == num_bytes:\n",
      "                return data\n",
      "            else:\n",
      "                read_data = self.read(num_bytes - l)\n",
      "            if read_data == '':\n",
      "                raise IOError('Could not read enough data.'\n",
      "                              '  Wanted %d bytes, got %d.' % (num_bytes, l))\n",
      "            data += read_data\n",
      "\n",
      "    def _read_check(self):\n",
      "        return struct.unpack(self.ENDIAN+self.HEADER_PREC,\n",
      "                             self._read_exactly(self._header_length)\n",
      "                            )[0]\n",
      "\n",
      "    def _write_check(self, number_of_bytes):\n",
      "        \"\"\"Write the header for the given number of bytes\"\"\"\n",
      "        self.write(struct.pack(self.ENDIAN+self.HEADER_PREC,\n",
      "                               number_of_bytes))\n",
      "\n",
      "    def readRecord(self):\n",
      "        \"\"\"Read a single fortran record\"\"\"\n",
      "        l = self._read_check()\n",
      "        data_str = self._read_exactly(l)\n",
      "        check_size = self._read_check()\n",
      "        if check_size != l:\n",
      "            raise IOError('Error reading record from data file')\n",
      "        return data_str\n",
      "\n",
      "    def writeRecord(self,s):\n",
      "        \"\"\"Write a record with the given bytes.\n",
      "\n",
      "        Parameters\n",
      "        ----------\n",
      "        s : the string to write\n",
      "\n",
      "        \"\"\"\n",
      "        length_bytes = len(s)\n",
      "        self._write_check(length_bytes)\n",
      "        self.write(s)\n",
      "        self._write_check(length_bytes)\n",
      "\n",
      "    def readString(self):\n",
      "        \"\"\"Read a string.\"\"\"\n",
      "        return self.readRecord()\n",
      "\n",
      "    def writeString(self,s):\n",
      "        \"\"\"Write a string\n",
      "\n",
      "        Parameters\n",
      "        ----------\n",
      "        s : the string to write\n",
      "        \n",
      "        \"\"\"\n",
      "        self.writeRecord(s)\n",
      "\n",
      "    _real_precisions = 'df'\n",
      "\n",
      "    def readReals(self, prec='f'):\n",
      "        \"\"\"Read in an array of real numbers.\n",
      "        \n",
      "        Parameters\n",
      "        ----------\n",
      "        prec : character, optional\n",
      "            Specify the precision of the array using character codes from\n",
      "            Python's struct module.  Possible values are 'd' and 'f'.\n",
      "            \n",
      "        \"\"\"\n",
      "        \n",
      "        _numpy_precisions = {'d': numpy.float64,\n",
      "                             'f': numpy.float32\n",
      "                            }\n",
      "\n",
      "        if prec not in self._real_precisions:\n",
      "            raise ValueError('Not an appropriate precision')\n",
      "            \n",
      "        data_str = self.readRecord()\n",
      "        num = len(data_str)/struct.calcsize(prec)\n",
      "        numbers =struct.unpack(self.ENDIAN+str(num)+prec,data_str) \n",
      "        return numpy.array(numbers, dtype=_numpy_precisions[prec])\n",
      "\n",
      "    def writeReals(self, reals, prec='f'):\n",
      "        \"\"\"Write an array of floats in given precision\n",
      "\n",
      "        Parameters\n",
      "        ----------\n",
      "        reals : array\n",
      "            Data to write\n",
      "        prec` : string\n",
      "            Character code for the precision to use in writing\n",
      "        \"\"\"\n",
      "        if prec not in self._real_precisions:\n",
      "            raise ValueError('Not an appropriate precision')\n",
      "        \n",
      "        # Don't use writeRecord to avoid having to form a\n",
      "        # string as large as the array of numbers\n",
      "        length_bytes = len(reals)*struct.calcsize(prec)\n",
      "        self._write_check(length_bytes)\n",
      "        _fmt = self.ENDIAN + prec\n",
      "        for r in reals:\n",
      "            self.write(struct.pack(_fmt,r))\n",
      "        self._write_check(length_bytes)\n",
      "    \n",
      "    _int_precisions = 'hilq'\n",
      "\n",
      "    def readInts(self, prec='i'):\n",
      "        \"\"\"Read an array of integers.\n",
      "        \n",
      "        Parameters\n",
      "        ----------\n",
      "        prec : character, optional\n",
      "            Specify the precision of the data to be read using \n",
      "            character codes from Python's struct module.  Possible\n",
      "            values are 'h', 'i', 'l' and 'q'\n",
      "            \n",
      "        \"\"\"\n",
      "        if prec not in self._int_precisions:\n",
      "            raise ValueError('Not an appropriate precision')\n",
      "            \n",
      "        data_str = self.readRecord()\n",
      "        num = len(data_str)/struct.calcsize(prec)\n",
      "        return numpy.array(struct.unpack(self.ENDIAN+str(num)+prec,data_str))\n",
      "\n",
      "    def writeInts(self, ints, prec='i'):\n",
      "        \"\"\"Write an array of integers in given precision\n",
      "\n",
      "        Parameters\n",
      "        ----------\n",
      "        reals : array\n",
      "            Data to write\n",
      "        prec : string\n",
      "            Character code for the precision to use in writing\n",
      "        \"\"\"\n",
      "        if prec not in self._int_precisions:\n",
      "            raise ValueError('Not an appropriate precision')\n",
      "        \n",
      "        # Don't use writeRecord to avoid having to form a\n",
      "        # string as large as the array of numbers\n",
      "        length_bytes = len(ints)*struct.calcsize(prec)\n",
      "        self._write_check(length_bytes)\n",
      "        _fmt = self.ENDIAN + prec\n",
      "        for item in ints:\n",
      "            self.write(struct.pack(_fmt,item))\n",
      "        self._write_check(length_bytes)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    }
   ],
   "metadata": {}
  }
 ]
}